# backend/agents/coordinator.py
"""
协调智能体 - 负责任务分配和响应整合
"""

from typing import List, Dict, Any
from langchain.schema import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from ..models.base_model import AIModelResponse
import json
import logging

logger = logging.getLogger(__name__)

class CoordinatorAgent:
    """协调智能体"""
    
    def __init__(self):
        self.name = "Coordinator"
        # 使用GPT-3.5作为协调模型（成本较低）
        self.coordinator_llm = None
    
    async def initialize(self, api_key: str):
        """初始化协调模型"""
        try:
            self.coordinator_llm = ChatOpenAI(
                model="gpt-3.5-turbo",
                openai_api_key=api_key,
                temperature=0.3
            )
            logger.info("协调智能体初始化成功")
        except Exception as e:
            logger.error(f"协调智能体初始化失败: {e}")
            raise
    
    async def analyze_request(
        self, 
        user_input: str, 
        available_models: List[str]
    ) -> Dict[str, Any]:
        """分析用户请求，决定使用哪些模型"""
        try:
            analysis_prompt = f"""
            作为AI协调员，分析用户请求并选择最适合的AI模型。
            
            用户请求: {user_input}
            可用模型: {', '.join(available_models)}
            
            模型特点:
            - gpt-4: 复杂推理、创意写作、深度分析
            - gpt-3.5-turbo: 快速响应、日常对话、简单任务
            - claude-3: 长文本处理、道德推理、谨慎回答
            
            请以JSON格式返回分析结果:
            {{
                "request_type": "问题类型",
                "complexity": "简单/中等/复杂",
                "selected_models": ["模型1", "模型2"],
                "reasoning": "选择理由",
                "metadata": {{"category": "分类", "priority": "优先级"}}
            }}
            """
            
            if self.coordinator_llm:
                messages = [
                    SystemMessage(content="你是一个AI模型协调专家，负责为用户请求选择最合适的AI模型。"),
                    HumanMessage(content=analysis_prompt)
                ]
                
                response = await self.coordinator_llm.ainvoke(messages)
                analysis = json.loads(response.content)
                
                logger.info(f"协调分析完成: {analysis['request_type']}")
                return analysis
            else:
                # 降级策略：简单的关键词匹配
                return self._fallback_analysis(user_input, available_models)
                
        except Exception as e:
            logger.error(f"请求分析失败: {e}")
            return self._fallback_analysis(user_input, available_models)
    
    def _fallback_analysis(self, user_input: str, available_models: List[str]) -> Dict[str, Any]:
        """降级分析策略"""
        keywords_complex = ["分析", "解释", "设计", "创作", "策略", "方案"]
        keywords_simple = ["什么", "怎么", "介绍", "列出"]
        
        is_complex = any(keyword in user_input for keyword in keywords_complex)
        is_simple = any(keyword in user_input for keyword in keywords_simple)
        
        if is_complex:
            selected = [m for m in ["gpt-4", "claude-3"] if m in available_models][:2]
            complexity = "复杂"
        elif is_simple:
            selected = [m for m in ["gpt-3.5-turbo", "gpt-4"] if m in available_models][:1]
            complexity = "简单"
        else:
            selected = available_models[:2]
            complexity = "中等"
        
        return {
            "request_type": "通用问题",
            "complexity": complexity,
            "selected_models": selected,
            "reasoning": "基于关键词的简单分析",
            "metadata": {"category": "general", "priority": "normal"}
        }
    
    async def synthesize_responses(
        self, 
        user_input: str, 
        responses: Dict[str, AIModelResponse]
    ) -> str:
        """整合多个AI的响应"""
        try:
            if len(responses) <= 1:
                return next(iter(responses.values())).content
            
            synthesis_prompt = f"""
            用户问题: {user_input}
            
            以下是不同AI模型的回答:
            """
            
            for model_name, response in responses.items():
                synthesis_prompt += f"\n【{model_name}的回答】:\n{response.content}\n"
            
            synthesis_prompt += """
            请整合这些回答，提供一个综合、准确、有用的最终回答。要求:
            1. 保留各回答的优点
            2. 去除重复内容
            3. 确保逻辑清晰
            4. 突出重点信息
            5. 如果回答有冲突，请指出并给出平衡的观点
            """
            
            if self.coordinator_llm:
                messages = [
                    SystemMessage(content="你是一个AI响应整合专家，擅长综合多个AI的回答。"),
                    HumanMessage(content=synthesis_prompt)
                ]
                
                response = await self.coordinator_llm.ainvoke(messages)
                logger.info("响应整合完成")
                return response.content
            else:
                # 降级策略：简单拼接
                return self._fallback_synthesis(responses)
                
        except Exception as e:
            logger.error(f"响应整合失败: {e}")
            return self._fallback_synthesis(responses)
    
    def _fallback_synthesis(self, responses: Dict[str, AIModelResponse]) -> str:
        """降级整合策略"""
        if not responses:
            return "抱歉，没有获得有效响应。"
        
        # 选择最长的响应作为主要答案
        best_response = max(responses.values(), key=lambda r: len(r.content))
        
        synthesis = f"**综合回答** (主要基于 {best_response.model_name}):\n\n"
        synthesis += best_response.content
        
        if len(responses) > 1:
            synthesis += f"\n\n**其他AI观点**:\n"
            for model_name, response in responses.items():
                if model_name != best_response.model_name:
                    preview = response.content[:100] + "..." if len(response.content) > 100 else response.content
                    synthesis += f"- **{model_name}**: {preview}\n"
        
        return synthesis

